package embedding

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
)

// OllamaConfig 表示Ollama服务的配置
type OllamaConfig struct {
	Host  string
	Port  int
	Model string
}

// OllamaEmbedder 表示Ollama嵌入服务
type OllamaEmbedder struct {
	config OllamaConfig
}

// EmbedRequest 表示嵌入请求
type EmbedRequest struct {
	Model  string `json:"model"`
	Prompt string `json:"prompt"`
}

// EmbedResponse 表示嵌入响应
type EmbedResponse struct {
	Embedding []float32 `json:"embedding"`
	Error     string    `json:"error,omitempty"`
}

// NewOllamaEmbedder 创建新的Ollama嵌入服务
func NewOllamaEmbedder(config OllamaConfig) *OllamaEmbedder {
	if config.Model == "" {
		// 默认使用bge-m3嵌入模型
		config.Model = "bge-m3"
	}

	if config.Host == "" {
		config.Host = "localhost"
	}

	if config.Port == 0 {
		config.Port = 11434
	}

	return &OllamaEmbedder{
		config: config,
	}
}

// GetEmbedding 获取文本的向量嵌入
func (e *OllamaEmbedder) GetEmbedding(text string) ([]float32, error) {
	url := fmt.Sprintf("http://%s:%d/api/embeddings", e.config.Host, e.config.Port)

	requestBody := EmbedRequest{
		Model:  e.config.Model,
		Prompt: text,
	}

	jsonData, err := json.Marshal(requestBody)
	if err != nil {
		return nil, fmt.Errorf("序列化请求失败: %w", err)
	}

	resp, err := http.Post(url, "application/json", bytes.NewBuffer(jsonData))
	if err != nil {
		return nil, fmt.Errorf("连接Ollama服务失败: %w", err)
	}
	defer resp.Body.Close()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("读取响应失败: %w", err)
	}

	var embedResponse EmbedResponse
	if err := json.Unmarshal(body, &embedResponse); err != nil {
		return nil, fmt.Errorf("解析响应失败: %w", err)
	}

	if embedResponse.Error != "" {
		return nil, fmt.Errorf("Ollama服务返回错误: %s", embedResponse.Error)
	}

	return embedResponse.Embedding, nil
}

// Embedder 表示嵌入接口
type Embedder interface {
	GetEmbedding(text string) ([]float32, error)
}
